#!/usr/bin/python

import logging
import random

from hips.scheduling import queue


class InitialConditionsMinConflictsSolver(object):
  """Minimum conflicts theory solver that chooses optimal initial conditions.

  Use the same principle of minimum conflicts, but we do a lot of additional
  pruning to get the search space down to a reasonable size before we start
  to solve.
  We take the following steps when looking for an initial assignment.
  1. Iterate through all the shifts and assign personnel based on selection
     via priority queue. The PQ uses the following criteria as heuristics
     on finding an initial arrangement:
     - Number of shifts already assigned
     - Surrounding values for the the shifts in question
     - Gives priority to satisfying night shifts
  2. If we could not find an initial assignment, go back and try again.
  3. After initial assignment, run forward checks using all constraints to
     further prune search space. (Right now we only do this to narrow down
     the "two consecutive shifts" criterion.)
  4. Solve using regular minimum conflicts method, with the exception that
     we only modify day shifts. This makes it easier to find a satisfying
     assignment because night shifts are hard to satisfy and we already have
     a possible arrangement from step 1. So now we are only really looking
     at the following criteria:
     - Max/min number of total shifts
     - Max number of consecutive shifts
  """

  def __init__(self, doctors, schedule, steps=1000):
    self._doctors = doctors
    self._schedule = schedule
    self._steps = steps
    self._iterations = 0

  def _AssignInitialCondition(self, domains):
    assignments = {}

    # Create priority queue for recording assignments.
    max_shifts = []
    max_night_shifts = []
    for doc in self._doctors:
      max_shifts.append(doc.max_shifts)
      max_night_shifts.append(doc.max_night_shifts)
    pq = queue.AssignmentPriorityQueue(len(self._doctors), max_shifts,
        max_night_shifts, self._schedule)

    # Note: Don't permute the days because this defeats the purpose
    # of the stateful priority queue.
    # Iterates through shifts in order.
    for i in xrange(len(domains)):
      excludes = []
      day = self._schedule.DayForShiftNumber(i)
      # Calculate exclusion list.
      for j, doc in enumerate(self._doctors):
        if day in doc.days_off:
          excludes.append(j)
      # TODO Inclusion list?
      if self._schedule.IsDay(i):  # TODO Generalize
        needed_for_shift = 2
      else:
        needed_for_shift = 1
      next = pq.NextAvailable(i, needed_for_shift, excludes)
      if not next:
        # TODO Better error handling
        # Warn that there are not enough docs for that day
        logging.info("Not enough doctors available for day %d shift %d" %
            (day, i))
        return None
      # Assign based on priority queue's suggestions.
      mask = 0
      for doc in next:
        mask |= 1 << doc
      assignments[i] = mask

    return assignments

  def _PushDomains(self, domains):
    for d in domains.values():
      d.PushState()

  def _PopDomains(self, domains):
    for d in domains.values():
      d.PopState()

  def _AssignAndGetSolution(self, domains, constraints, vconstraints,
      max_local_iterations):
    """Performs min conflicts solver with a few important changes.

    Optimizes the initial arrangement for a schedule before starting
    search for a solution. This initial heuristic places us much closer
    to a solution than a completely random initial assignment.
    """
    # Attempt initial assignment as many times as required.
    # This generates a new random assignment each time.
    self._iterations += 1
    assignments = self._AssignInitialCondition(domains)
    if not assignments:
      logging.info("Didn't find a reasonable assignment. "
          "Iterations: %d" % self._iterations)
      return None

    logging.info("Starting with initial assignment:")
    logging.info(str(assignments))

    # Forward check to prune as much as possible before attempting search.
    self._PushDomains(domains)
    lst = domains.keys()
    for variable in lst:
      for constraint, variables in vconstraints[variable]:
        if not constraint.ForwardCheck(variables, domains, assignments):
          logging.info("Forward check shows no solutions for assignment. "
              "Iterations: %d" % self._iterations)
          return None

    local_iterations = 0
    conflicts = 0

    while (local_iterations < max_local_iterations and
        self._iterations < self._steps):
      self._iterations += 1
      local_iterations += 1
      conflicts = 0
      mincount = 0
      lst = domains.keys()
      random.shuffle(lst)

      for variable in lst:
        # Check if variable is not in conflict
        conflicts = 0
        for constraint, variables in vconstraints[variable]:
          conflicts += constraint(variables, domains, assignments)
        # No conflicts, break the search loop and return
        if conflicts == 0:
          continue

        # Variable has conflicts. Find values with less conflicts.
        mincount = conflicts
        minvalues = []
        for value in domains[variable]:
          assignments[variable] = value
          count = 0
          for constraint, variables in vconstraints[variable]:
            count += constraint(variables, domains, assignments)
          if count == mincount:
            minvalues.append(value)
          elif count < mincount:
            mincount = count
            del minvalues[:]
            minvalues.append(value)

        # Pick a random one from these values.
        if minvalues:
          assignments[variable] = random.choice(minvalues)

      logging.info("%d conflicts in current assignment" % conflicts)
      if mincount < conflicts:
        logging.info("  --> min value has %d conflicts" % mincount)

      if not conflicts:
        # No need to pop here because we don't need state any more.
        # We're done with the search and all the solving.
        return assignments

    self._PopDomains(domains)
    if conflicts > 0:
      logging.info("WARNING: Assignment has %d conflicts." % conflicts)
    return assignments

  def Iterations(self):
    return self._iterations

  def Solve(self, domains, constraints, vconstraints):
    self._iterations = 0
    while self._iterations < self._steps:
      if self._iterations:
        logging.info("Generating new assignment. Iterations: %d" %
            self._iterations)
      # The last parameter is local_iterations. It is the same as _steps
      # because the fine-grained min conflicts solver will converge to an
      # answer with as few conflicts as possible and that answer may still
      # be useful.
      assignments = self._AssignAndGetSolution(domains, constraints,
          vconstraints, self._steps)
      if assignments:
        return assignments
    return None
